﻿///
/// @file ILocArm64.cpp
/// @brief 指令序列管理的实现，ILOC的全称为Intermediate Language for Optimizing Compilers
/// @author zenglj (zenglj@live.com)
/// @version 1.0
/// @date 2024-11-21
///
/// @copyright Copyright (c) 2024
///
/// @par 修改日志:
/// <table>
/// <tr><th>Date       <th>Version <th>Author  <th>Description
/// <tr><td>2024-11-21 <td>1.0     <td>zenglj  <td>新做
/// </table>
///
#include <cstdio>
#include <string>
#include <cstring>

#include "ILocArm64.h"
#include "Common.h"
#include "Function.h"
#include "PlatformArm64.h"
#include "Module.h"
#include "ConstInt.h"
#include "ConstFloat.h"
#include "SimpleRegisterAllocator.h"

ArmInst::ArmInst(std::string _opcode,
                 std::string _result,
                 std::string _arg1,
                 std::string _arg2,
                 std::string _cond,
                 std::string _addition)
    : opcode(_opcode), cond(_cond), result(_result), arg1(_arg1), arg2(_arg2), addition(_addition), dead(false)
{}

/*
    指令内容替换
*/
void ArmInst::replace(std::string _opcode,
                      std::string _result,
                      std::string _arg1,
                      std::string _arg2,
                      std::string _cond,
                      std::string _addition)
{
    opcode = _opcode;
    result = _result;
    arg1 = _arg1;
    arg2 = _arg2;
    cond = _cond;
    addition = _addition;

#if 0
    // 空操作，则设置为dead
    if (op == "") {
        dead = true;
    }
#endif
}

/*
    设置为无效指令
*/
void ArmInst::setDead()
{
    dead = true;
}

/*
    输出函数
*/
std::string ArmInst::outPut()
{
    // 无用代码，什么都不输出
    if (dead) {
        return "";
    }

    // 占位指令,可能需要输出一个空操作，看是否支持 FIXME
    if (opcode.empty()) {
        return "";
    }

    std::string ret = opcode;

    if (!cond.empty()) {
        ret += cond;
    }

    // 结果输出
    if (!result.empty()) {
        if (result == ":") {
            ret += result;
        } else {
            ret += " " + result;
        }
    }

    // 第一元参数输出
    if (!arg1.empty()) {
        ret += "," + arg1;
    }

    // 第二元参数输出
    if (!arg2.empty()) {
        ret += "," + arg2;
    }

    // 其他附加信息输出
    if (!addition.empty()) {
        ret += "," + addition;
    }

    return ret;
}

#define emit(...) code.push_back(new ArmInst(__VA_ARGS__))

// 新增：获取64位寄存器名
static std::string getRegName64(int reg_no)
{
    return PlatformArm64::regName[reg_no];
}

/// @brief 构造函数
/// @param _module 符号表
ILocArm64::ILocArm64(Module * _module)
{
    this->module = _module;
}

/// @brief 析构函数
ILocArm64::~ILocArm64()
{
    std::list<ArmInst *>::iterator pIter;

    for (pIter = code.begin(); pIter != code.end(); ++pIter) {
        delete (*pIter);
    }
}

/// @brief 删除无用的Label指令
void ILocArm64::deleteUsedLabel()
{
    std::list<ArmInst *> labelInsts;
    for (ArmInst * arm: code) {
        if ((!arm->dead) && (arm->opcode[0] == '.') && (arm->result == ":")) {
            labelInsts.push_back(arm);
        }
    }

    for (ArmInst * labelArm: labelInsts) {
        bool labelUsed = false;

        for (ArmInst * arm: code) {
            // TODO 转移语句的指令标识符根据定义修改判断
            if ((!arm->dead) && (arm->opcode[0] == 'b') && (arm->result == labelArm->opcode)) {
                labelUsed = true;
                break;
            }
        }

        if (!labelUsed) {
            labelArm->setDead();
        }
    }
}

/// @brief 输出汇编
/// @param file 输出的文件指针
/// @param outputEmpty 是否输出空语句
void ILocArm64::outPut(FILE * file, bool outputEmpty)
{
    for (auto arm: code) {

        std::string s = arm->outPut();

        if (arm->result == ":") {
            // Label指令，不需要Tab输出
            fprintf(file, "%s\n", s.c_str());
            continue;
        }

        if (!s.empty()) {
            fprintf(file, "\t%s\n", s.c_str());
        } else if ((outputEmpty)) {
            fprintf(file, "\n");
        }
    }
}

/// @brief 获取当前的代码序列
/// @return 代码序列
std::list<ArmInst *> & ILocArm64::getCode()
{
    return code;
}

/**
 * 数字变字符串，若flag为真，则变为立即数寻址（加#）
 */
std::string ILocArm64::toStr(int num, bool flag)
{
    std::string ret;

    if (flag) {
        ret = "#";
    }

    ret += std::to_string(num);

    return ret;
}

/*
    产生标签
*/
void ILocArm64::label(std::string name)
{
    // .L1:
    emit(name, ":");
}

/// @brief 0个源操作数指令
/// @param op 操作码
/// @param rs 操作数
void ILocArm64::inst(std::string op, std::string rs)
{
    emit(op, rs);
}

/// @brief 一个操作数指令
/// @param op 操作码
/// @param rs 操作数
/// @param arg1 源操作数
void ILocArm64::inst(std::string op, std::string rs, std::string arg1)
{
    emit(op, rs, arg1);
}

/// @brief 一个操作数指令
/// @param op 操作码
/// @param rs 操作数
/// @param arg1 源操作数
/// @param arg2 源操作数
void ILocArm64::inst(std::string op, std::string rs, std::string arg1, std::string arg2)
{
    emit(op, rs, arg1, arg2);
}

///
/// @brief 注释指令，不包含分号
///
void ILocArm64::comment(std::string str)
{
    emit("//", str);
}

/*
    加载立即数 movz/movk 组合加载64位立即数
*/
void ILocArm64::load_imm(int rs_reg_no, int64_t constant)
{
    std::string reg_name = PlatformArm64::regName[rs_reg_no];

    // 处理负数
    if (constant < 0) {
        // 对于负数，使用补码表示
        uint64_t uconst = static_cast<uint64_t>(constant);

        // 分解为4个16位段
        uint16_t w0 = uconst & 0xFFFF;
        uint16_t w1 = (uconst >> 16) & 0xFFFF;
        uint16_t w2 = (uconst >> 32) & 0xFFFF;
        uint16_t w3 = (uconst >> 48) & 0xFFFF;

        // 使用movz初始化最低有效的非零段
        bool first = true;
        if (w0 != 0) {
            emit("movz", reg_name, "#" + std::to_string(w0));
            first = false;
        }
        if (w1 != 0) {
            if (first) {
                emit("movz", reg_name, "#" + std::to_string(w1), "lsl #16");
                first = false;
            } else {
                emit("movk", reg_name, "#" + std::to_string(w1), "lsl #16");
            }
        }
        if (w2 != 0) {
            if (first) {
                emit("movz", reg_name, "#" + std::to_string(w2), "lsl #32");
                first = false;
            } else {
                emit("movk", reg_name, "#" + std::to_string(w2), "lsl #32");
            }
        }
        if (w3 != 0) {
            if (first) {
                emit("movz", reg_name, "#" + std::to_string(w3), "lsl #48");
            } else {
                emit("movk", reg_name, "#" + std::to_string(w3), "lsl #48");
            }
        }

        if (first) {
            // 所有段都为0，但这是负数，说明是0x0000000000000000的情况，实际上不应该到这里
            // 对于-1等情况，至少有一个段不为0
            emit("movz", reg_name, "#0");
        }
    }
    // 处理 [0, 65536) 范围内的正数
    else if (constant >= 0 && constant < 65536) {
        emit("movz", reg_name, "#" + std::to_string(constant));
    }
    // 处理超出 [0, 65536) 范围的正数
    else {
        uint64_t uconst = static_cast<uint64_t>(constant);

        // 分解为4个16位段
        uint16_t w0 = uconst & 0xFFFF;
        uint16_t w1 = (uconst >> 16) & 0xFFFF;
        uint16_t w2 = (uconst >> 32) & 0xFFFF;
        uint16_t w3 = (uconst >> 48) & 0xFFFF;

        // 使用movz初始化最低有效的非零段
        bool first = true;
        if (w0 != 0) {
            emit("movz", reg_name, "#" + std::to_string(w0));
            first = false;
        }
        if (w1 != 0) {
            if (first) {
                emit("movz", reg_name, "#" + std::to_string(w1), "lsl #16");
                first = false;
            } else {
                emit("movk", reg_name, "#" + std::to_string(w1), "lsl #16");
            }
        }
        if (w2 != 0) {
            if (first) {
                emit("movz", reg_name, "#" + std::to_string(w2), "lsl #32");
                first = false;
            } else {
                emit("movk", reg_name, "#" + std::to_string(w2), "lsl #32");
            }
        }
        if (w3 != 0) {
            if (first) {
                emit("movz", reg_name, "#" + std::to_string(w3), "lsl #48");
            } else {
                emit("movk", reg_name, "#" + std::to_string(w3), "lsl #48");
            }
        }

        if (first) {
            // 所有段都为0，直接使用movz加载0
            emit("movz", reg_name, "#0");
        }
    }
}

void ILocArm64::load_imm_float(int rs_reg_no, float constant)
{
    // 1. 将 float 常量按 IEEE 754 位模式解释为 uint32_t
    uint32_t bits;
    std::memcpy(&bits, &constant, sizeof(bits));

    // 2. 生成指令：将 bits 拆成 movz + movk 的组合加载到通用寄存器（如 wTmp）
    int wTmp = ARM64_TMP_REG_NO;

    uint16_t lower16 = bits & 0xFFFF;
    uint16_t mid16 = (bits >> 16) & 0xFFFF;

    if (lower16 != 0) {
        emit("movz", PlatformArm64::regNameW[wTmp], "#" + std::to_string(lower16));
        if (mid16 != 0) {
            emit("movk", PlatformArm64::regNameW[wTmp], "#" + std::to_string(mid16), "lsl #16");
        }
    } else if (mid16 != 0) {
        emit("movz", PlatformArm64::regNameW[wTmp], "#" + std::to_string(mid16), "lsl #16");
    } else {
        emit("movz", PlatformArm64::regNameW[wTmp], "#0");
    }

    // 3. 将整数位模式 move 到浮点寄存器 s[rs_reg_no]
    emit("fmov", PlatformArm64::regNameS[rs_reg_no], PlatformArm64::regNameW[wTmp]);
}

/// @brief 加载符号值 ldr r0,=g; ldr r0,[r0]
/// @param rs_reg_no 结果寄存器编号
/// @param name 符号名
void ILocArm64::load_symbol(int rs_reg_no, std::string name)
{
    // ARM64: adrp/add 组合加载符号地址，简化为adrp
    emit("adrp", PlatformArm64::regName[rs_reg_no], name);
    emit("add", PlatformArm64::regName[rs_reg_no], PlatformArm64::regName[rs_reg_no], ":lo12:" + name);
}

/// @brief 基址寻址 ldr r0,[fp,#100]
/// @param rsReg 结果寄存器
/// @param base_reg_no 基址寄存器
/// @param offset 偏移
void ILocArm64::load_base(int rs_reg_no, int base_reg_no, int offset)
{
    std::string rsReg = PlatformArm64::regName[rs_reg_no];
    std::string base = PlatformArm64::regName[base_reg_no];
    if (PlatformArm64::isDisp(offset)) {
        if (offset)
            base += "," + toStr(offset);
    } else {
        load_imm(rs_reg_no, offset);
        base += "," + rsReg;
    }
    base = "[" + base + "]";
    emit("ldr", rsReg, base);
}

// 从基址寄存器加偏移（或加寄存器）地址处加载内存到寄存器。
void ILocArm64::load_base_i(int rs_reg_no, int base_reg_no, int disp)
{
    std::string rsReg = PlatformArm64::regName[rs_reg_no];
    // 检查base_reg_no是否有效
    if (base_reg_no < 0) {
        // 如果base_reg_no无效，直接加载立即数
        load_imm(rs_reg_no, disp);
        return;
    }

    std::string base = PlatformArm64::regName[base_reg_no];
    if (PlatformArm64::isDisp(disp)) {
        if (disp)
            base += "," + toStr(disp);
    } else {
        load_imm(rs_reg_no, disp);
        base += "," + rsReg;
    }
    base = "[" + base + "]";
    emit("ldr", rsReg, base);
}
// 从基址寄存器加偏移（或加寄存器）地址处加载内存到寄存器。
void ILocArm64::load_base_f(int rs_reg_no, int base_reg_no, int disp)
{
    std::string rsReg = PlatformArm64::regNameS[rs_reg_no];
    std::string base = PlatformArm64::regName[base_reg_no];
    std::string tmp = PlatformArm64::regName[ARM64_TMP_REG_NO];
    if (PlatformArm64::isDisp(disp)) {
        if (disp)
            base += "," + toStr(disp);
    } else {
        load_imm(ARM64_TMP_REG_NO, disp);
        base += "," + tmp;
    }
    base = "[" + base + "]";
    emit("ldr", rsReg, base);
}

/// @brief 加载数组基址（带偏移）
/// @param rs_reg_no 结果寄存器编号
/// @param base_reg_no 基址寄存器编号
/// @param disp 偏移量
void ILocArm64::load_array_base(int rs_reg_no, int base_reg_no, int disp, bool is_param)
{
    std::string rsReg = PlatformArm64::regName[rs_reg_no];
    std::string base = PlatformArm64::regName[base_reg_no];
    if (PlatformArm64::isDisp(disp)) {
        base += "," + toStr(disp);
    } else {
        load_imm(rs_reg_no, disp);
        base += "," + rsReg;
    }
    emit("add", rsReg, base);
    if (is_param) {
        // !如果是参数传递，使用ldr提取地址
        emit("ldr", rsReg, "[" + rsReg + "]");
    }
}

/// @brief 基址寻址 str r0,[fp,#100]
/// @param srcReg 源寄存器
/// @param base_reg_no 基址寄存器
/// @param disp 偏移
/// @param tmp_reg_no 可能需要临时寄存器编号
void ILocArm64::store_base(int src_reg_no, int base_reg_no, int disp, int tmp_reg_no)
{
    std::string base = getRegName64(base_reg_no);

    if (PlatformArm64::isDisp(disp)) {
        // 有效的偏移常量

        // 若disp为0，则直接采用基址，否则采用基址+偏移
        // [fp,#-16] [fp]
        if (disp) {
            base += "," + toStr(disp);
        }
    } else {
        // 先把立即数赋值给指定的寄存器tmpReg，然后采用基址+寄存器的方式进行

        // ldr r9,=-4096
        load_imm(tmp_reg_no, disp);

        // fp,r9
        base += "," + getRegName64(tmp_reg_no);
    }

    // 内存间接寻址
    base = "[" + base + "]";

    // str r8,[fp,#-16]
    // str r8,[fp,r9]
    emit("str", getRegName64(src_reg_no), base);
}

// 将寄存器内容存储到基址寄存器加偏移（或加寄存器）地址处。
void ILocArm64::store_base_i(int src_reg_no, int base_reg_no, int disp, int tmp_reg_no)
{
    std::string base = PlatformArm64::regName[base_reg_no];
    if (PlatformArm64::isDisp(disp)) {
        if (disp)
            base += "," + toStr(disp);
    } else {
        load_imm(tmp_reg_no, disp);
        base += "," + PlatformArm64::regName[tmp_reg_no];
    }
    base = "[" + base + "]";
    emit("str", PlatformArm64::regName[src_reg_no], base);
}

void ILocArm64::store_base_f(int src_reg_no, int base_reg_no, int disp, int tmp_reg_no)
{
    std::string base = PlatformArm64::regName[base_reg_no];
    if (PlatformArm64::isDisp(disp)) {
        if (disp)
            base += "," + toStr(disp);
    } else {
        load_imm(tmp_reg_no, disp);
        base += "," + PlatformArm64::regName[tmp_reg_no];
    }
    base = "[" + base + "]";
    emit("str", PlatformArm64::regNameS[src_reg_no], base);
}

/// @brief 寄存器Mov操作
/// @param rs_reg_no 结果寄存器
/// @param src_reg_no 源寄存器
void ILocArm64::mov_reg(int rs_reg_no, int src_reg_no)
{
    emit("mov", getRegName64(rs_reg_no), getRegName64(src_reg_no));
}

/// @brief 加载变量到寄存器，保证将变量放到reg中
/// @param rs_reg_no 结果寄存器
/// @param src_var 源操作数
/// @param is_float_var 是否为浮点变量
/// @param is_param 是否为参数
void ILocArm64::load_var(int rs_reg_no, Value * src_var, bool is_float_var, bool is_param)
{
    // 添加空指针检查
    if (!src_var) {
        minic_log(LOG_ERROR, "Null variable in load_var");
        return;
    }

    if (Instanceof(constIntVal, ConstInt *, src_var)) {
        load_imm(rs_reg_no, constIntVal->getVal());
    } else if (Instanceof(constFloatVal, ConstFloat *, src_var)) {
        load_imm_float(rs_reg_no, constFloatVal->getVal());
    } else if (src_var->getRegId() != -1) {
        // 添加这个关键的检查！
        int32_t src_regId = src_var->getRegId();
        if (src_regId != rs_reg_no) {
            if (is_float_var) {
                emit("fmov", PlatformArm64::regNameS[rs_reg_no], PlatformArm64::regNameS[src_regId]);
            } else {
                emit("mov", PlatformArm64::regName[rs_reg_no], PlatformArm64::regName[src_regId]);
            }
        }
    } else if (Instanceof(globalVar, GlobalVariable *, src_var)) {
        // 全局变量处理
        load_symbol(rs_reg_no, globalVar->getName());

        // 检查变量类型
        Type * varType = src_var->getType();
        if (!varType) {
            minic_log(LOG_WARNING, "Global variable has no type information");
            emit("ldr", PlatformArm64::regName[rs_reg_no], "[" + PlatformArm64::regName[rs_reg_no] + "]");
        } else if (!varType->isArrayType()) {
            if (is_float_var) {
                emit("ldr", PlatformArm64::regNameS[rs_reg_no], "[" + PlatformArm64::regName[rs_reg_no] + "]");
            } else {
                emit("ldr", PlatformArm64::regName[rs_reg_no], "[" + PlatformArm64::regName[rs_reg_no] + "]");
            }
        }
    } else {
        // 局部变量，通过栈帧偏移加载
        int32_t var_baseRegId = -1;
        int64_t var_offset = -1;
        bool result = src_var->getMemoryAddr(&var_baseRegId, &var_offset);
        if (!result) {
            minic_log(LOG_ERROR, "Failed to get memory address for variable");
            return;
        }

        // 检查变量类型
        Type * varType = src_var->getType();
        if (varType && varType->isArrayType()) {
            load_array_base(rs_reg_no, var_baseRegId, var_offset, is_param);
        } else {
            if (is_float_var) {
                load_base_f(rs_reg_no, var_baseRegId, var_offset);
            } else {
                load_base_i(rs_reg_no, var_baseRegId, var_offset);
            }
        }
    }
}

/// @brief 加载变量地址到寄存器
/// @param rs_reg_no
/// @param var
void ILocArm64::lea_var(int rs_reg_no, Value * var)
{
    if (Instanceof(globalVar, GlobalVariable *, var)) {
        // 全局变量，直接加载符号地址
        load_symbol(rs_reg_no, globalVar->getName());
        return;
    }
    // 局部变量原有逻辑
    int32_t var_baseRegId = -1;
    int64_t var_offset = -1;
    bool result = var->getMemoryAddr(&var_baseRegId, &var_offset);
    if (!result) {
        minic_log(LOG_ERROR, "BUG");
    }
    leaStack(rs_reg_no, var_baseRegId, var_offset);
}

/// @brief 保存寄存器到变量，保证将计算结果（r8）保存到变量
/// @param src_reg_no 源寄存器
/// @param dest_var  变量
/// @param tmp_reg_no 第三方寄存器
void ILocArm64::store_var(int src_reg_no, Value * dest_var, int tmp_reg_no, bool is_float_var)
{
    // 添加空指针检查
    if (!dest_var) {
        minic_log(LOG_ERROR, "Null variable in store_var");
        return;
    }
    static SimpleRegisterAllocator simpleRegisterAllocator; // 添加静态寄存器分配器实例

    if (dest_var->getRegId() != -1) {
        int dest_reg_id = dest_var->getRegId();
        if (src_reg_no != dest_reg_id) {
            if (is_float_var) {
                // 如果是浮点类型，使用fmov
                emit("fmov", PlatformArm64::regNameS[dest_reg_id], PlatformArm64::regNameS[src_reg_no]);
            } else {

                emit("mov", PlatformArm64::regName[dest_reg_id], PlatformArm64::regName[src_reg_no]);
            }
        }
    } else if (Instanceof(globalVar, GlobalVariable *, dest_var)) {
        // 如果 src_reg_no == tmp_reg_no，会导致 str x9, [x9]这类错误
        if (src_reg_no == tmp_reg_no) {
            tmp_reg_no = simpleRegisterAllocator.Allocate(); // 重新申请临时寄存器
        }

        load_symbol(tmp_reg_no, globalVar->getName());

        if (is_float_var) {
            emit("str", PlatformArm64::regNameS[src_reg_no], "[" + PlatformArm64::regNameS[tmp_reg_no] + "]");
        } else {
            emit("str", PlatformArm64::regName[src_reg_no], "[" + PlatformArm64::regName[tmp_reg_no] + "]");
        }

        // 如果我们重新申请了寄存器，就释放它
        if (src_reg_no == tmp_reg_no) {
            simpleRegisterAllocator.free(tmp_reg_no);
        }
    } else {
        int32_t dest_baseRegId = -1;
        int64_t dest_offset = -1;
        bool result = dest_var->getMemoryAddr(&dest_baseRegId, &dest_offset);
        if (!result) {
            minic_log(LOG_ERROR, "BUG");
        }
        if (is_float_var) {
            // 如果是浮点类型，使用store_base_f
            store_base_f(src_reg_no, dest_baseRegId, dest_offset, tmp_reg_no);
        } else
            store_base_i(src_reg_no, dest_baseRegId, dest_offset, tmp_reg_no);
    }
}

/// @brief 加载栈内变量地址
/// @param rsReg 结果寄存器号
/// @param base_reg_no 基址寄存器
/// @param off 偏移
void ILocArm64::leaStack(int rs_reg_no, int base_reg_no, int off)
{
    std::string rs_reg_name = getRegName64(rs_reg_no);
    std::string base_reg_name = getRegName64(base_reg_no);

    if (PlatformArm64::constExpr(off))
        // add r8,fp,#-16
        emit("add", rs_reg_name, base_reg_name, toStr(off));
    else {
        // ldr r8,=-257
        load_imm(rs_reg_no, off);

        // add r8,fp,r8
        emit("add", rs_reg_name, base_reg_name, rs_reg_name);
    }
}

/// @brief 函数内栈内空间分配（局部变量、形参变量、函数参数传值，或不能寄存器分配的临时变量等）
/// @param func 函数
/// @param tmp_reg_No
void ILocArm64::allocStack(Function * func, int tmp_reg_no)
{
    int funcCallArgCnt = func->getMaxFuncCallArgCnt() - 8;
    funcCallArgCnt = std::max(funcCallArgCnt, 0);

    // 16字节对齐实参空间
    int arg_space = funcCallArgCnt * 8;
    arg_space += (16 - arg_space % 16) % 16;

    int off = func->getMaxDep();
    // 16字节对齐局部空间
    off += (16 - off % 16) % 16;

    off += arg_space;

    if (0 == off)
        return;

    if (PlatformArm64::constExpr(off)) {
        emit("sub", "sp", "sp", toStr(off));
    } else {
        load_imm(tmp_reg_no, off);
        emit("sub", "sp", "sp", PlatformArm64::regName[tmp_reg_no]);
    }

    // 16字节对齐后设置fp
    inst("add", PlatformArm64::regName[ARM64_FP_REG_NO], "sp", toStr(arg_space));
}

/// @brief 调用函数fun
/// @param fun
void ILocArm64::call_fun(std::string name)
{
    // 函数返回值在r0,不需要保护
    emit("bl", name);
}

/// @brief NOP操作
void ILocArm64::nop()
{
    // FIXME 无操作符，要确认是否用nop指令
    emit("");
}

///
/// @brief 无条件跳转指令
/// @param label 目标Label名称
///
void ILocArm64::jump(std::string label)
{
    emit("b", label);
}